package com.itcam.cryptography;

public class AESDemo {

    public static void main(String[] args) {
        short[] plaintext = {
//                0x32, 0x43, 0xf6, 0xa8,
//                0x88, 0x5a, 0x30, 0x8d,
//                0x31, 0x31, 0x98, 0xa2,
//                0xe0, 0x37, 0x07, 0x34

//                0x19, 0x3d, 0xe3, 0xbe,
//                0xa0, 0xf4, 0xe2, 0x2b,
//                0x9a, 0xc6, 0x8d, 0x2a,
//                0xe9, 0xf8, 0x48, 0x08, 213110085

                0x21, 0x31, 0x10, 0x08,
                0x50, 0x00, 0x00, 0x00,
                0x00, 0x00, 0x00, 0x00,
                0x00, 0x00, 0x00, 0x00
        };
        short[] key = {
//                0x2b, 0x7e, 0x15, 0x16,
//                0x28, 0xae, 0xd2, 0xa6,
//                0xab, 0xf7, 0x15, 0x88,
//                0x09, 0xcf, 0x4f, 0x3c

//                0xa0, 0xfa, 0xfe, 0x17,
//                0x88, 0x54, 0x2c, 0xb1,
//                0x23, 0xa3, 0x39, 0x39,
//                0x2a, 0x6c, 0x76, 0x05

                0xa0, 0xfa, 0xfe, 0x17,
                0x88, 0x54, 0x2c, 0xb1,
                0x23, 0xa3, 0x39, 0x39,
                0x2a, 0x6c, 0x76, 0x05

        };
        short[][] input = new short[4][4];
        short[][] output = new short[4][4];
        short[][] subkey = new short[4][4];


        System.out.println("明文为：");
        for (int i = 0; i < 16; i++) {
            System.out.printf("%02X ", plaintext[i]);
        }
        System.out.println();
        System.out.println();

        System.out.println("初始密钥为：");
        for (int i = 0; i < 16; i++) {
            System.out.printf("%02X ", key[i]);
        }
        System.out.println();
        System.out.println();
        System.out.println();

        // 生成状态矩阵
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                input[j][i] = plaintext[j * 4 + i];
            }
        }
        System.out.println("输入明文状态矩阵：");
        printstate(input);
        System.out.println();

        // 生成初始密钥
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                subkey[j][i] = key[j * 4 + i];
            }
        }
        System.out.println("输入128比特初始密钥矩阵:");
        printstate(subkey);
        System.out.println();

        AddRoundKey(input, subkey);
        System.out.println("初始密钥加运算输出：");
        printstate(input);
        System.out.println();

//        SubBytes(input);
//        System.out.println("字节替代输出：");
//        printstate(input);
//        System.out.println();

        ShiftRows(input);
        System.out.println("行移位输出：");
        printstate(input);
        System.out.println();


    }


    public static void printstate(short[][] array) {
        for (int i = 0; i < array.length; i++) {
            for (int j = 0; j < array[i].length; j++) {
                System.out.print(array[i][j] + " ");
            }
            System.out.println();
        }
    }

    public static void AddRoundKey(short[][] state, short[][] round) {

        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                state[i][j] ^= round[i][j];
            }
        }
    }

    public static void SubBytes(short[][] state) {
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                int row = (state[i][j] & 0xF0) >> 4;
                int col = state[i][j] & 0x0F;
                state[i][j] = (short) S_BOX[row * 16 + col];
            }
        }
    }

    public static void ShiftRows(short[][] state) {
        int[][] temp = new int[4][4];
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                temp[i][j] = state[i][(j + i) % 4];
            }
        }
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                state[i][j] = (short) temp[i][j];
            }
        }
    }


    private static final int[] S_BOX = {
            0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    };

    private static final int[] RCON = {
            0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36
    };

    private static final int[] ROW_SHIFT = {
            0, 1, 2, 3
    };

    private static final int[] COL_MIX = {
            0x02, 0x03, 0x01, 0x01
    };


}
